#ifndef TIEC_TREE_PATH_H
#define TIEC_TREE_PATH_H

#include <optional>
#include "tree.h"

namespace NS_TIEC {
  /// 语法树短路遍历器基类，用于在语法树中搜寻结果，只要搜寻到第一个结果就停止遍历
  template <typename Result, typename Param>
  class TreeSearcher : public TreeVisitor<std::optional<Result>, Param> {
  public:
    template <typename T, typename = std::enable_if_t<std::is_base_of_v<Tree, T>>>
    std::optional<Result> searchTree(const PTR<T>& tree, Param param) {
      return this->visitTree(tree, param);
    }

    template <typename T, typename = std::enable_if_t<std::is_base_of_v<Tree, T>>>
    std::optional<Result> searchTrees(const List<PTR<T>>& list, Param param) {
      for (const PTR<T>& tree : list) {
        std::optional<Result> r = searchTree(tree, param);
        if (r.has_value()) {
          return r;
        }
      }
      return std::nullopt;
    }

    std::optional<Result> visitCompilationUnitTree(const PTR<CompilationUnitTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->pkg, param);
      if (r.has_value()) return r;
      r = searchTrees<AnnotationTree>(tree->annotations, param);
      if (r.has_value()) return r;
      r = searchTrees<Tree>(tree->defs, param);
      if (r.has_value()) return r;
      return std::nullopt;
    }

    std::optional<Result> visitClassTree(const PTR<ClassTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->super, param);
      if (r.has_value()) return r;
      r = searchTrees<AnnotationTree>(tree->annotations, param);
      if (r.has_value()) return r;
      for (const PAIR<PTR<IdentifierTree>, PTR<ExpressionTree>>& pair : tree->templates) {
        r = searchTree(pair.first, param);
        if (r.has_value()) return r;
        r = searchTree(pair.second, param);
        if (r.has_value()) return r;
      }
      r = searchTrees<AnnotatedTree>(tree->members, param);
      if (r.has_value()) return r;
      r = searchTrees<CodeStatementTree>(tree->codes, param);
      if (r.has_value()) return r;
      return std::nullopt;
    }

    std::optional<Result> visitVariableTree(const PTR<VariableTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<AnnotationTree>(tree->annotations, param);
      if (r.has_value()) return r;
      r = searchTree(tree->type_expr, param);
      if (r.has_value()) return r;
      r = searchTree(tree->initializer, param);
      return r;
    }

    std::optional<Result> visitMethodTree(const PTR<MethodTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<AnnotationTree>(tree->annotations, param);
      if (r.has_value()) return r;
      r = searchTrees<VariableTree>(tree->parameters, param);
      if (r.has_value()) return r;
      r = searchTree(tree->return_type, param);
      if (r.has_value()) return r;
      r = searchTree(tree->body, param);
      return r;
    }

    std::optional<Result> visitEventTree(const PTR<EventTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<AnnotationTree>(tree->annotations, param);
      if (r.has_value()) return r;
      r = searchTree(tree->event_obj, param);
      if (r.has_value()) return r;
      r = searchTrees<VariableTree>(tree->parameters, param);
      if (r.has_value()) return r;
      r = searchTree(tree->return_type, param);
      if (r.has_value()) return r;
      r = searchTree(tree->body, param);
      return r;
    }

    std::optional<Result> visitAnnotationTree(const PTR<AnnotationTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<ExpressionTree>(tree->args, param);
      return r;
    }

    std::optional<Result> visitCodeStatementTree(const PTR<CodeStatementTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<CodeSegmentTree>(tree->lines, param);
      return r;
    }

    std::optional<Result> visitCodeExpressionTree(const PTR<CodeExpressionTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<CodeSegmentTree>(tree->lines, param);
      return r;
    }

    std::optional<Result> visitCodeSegmentTree(const PTR<CodeSegmentTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<ExpressionTree>(tree->trees, param);
      return r;
    }

    std::optional<Result> visitBlockTree(const PTR<BlockTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<StatementTree>(tree->statements, param);
      return r;
    }

    std::optional<Result> visitExpressionStatementTree(const PTR<ExpressionStatementTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->expr, param);
      return r;
    }

    std::optional<Result> visitIfTree(const PTR<IfTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->condition, param);
      if (r.has_value()) return r;
      r = searchTree(tree->then_part, param);
      if (r.has_value()) return r;
      r = searchTree(tree->else_part, param);
      return r;
    }

    std::optional<Result> visitSwitchTree(const PTR<SwitchTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->selector, param);
      if (r.has_value()) return r;
      r = searchTrees<CaseTree>(tree->cases, param);
      if (r.has_value()) return r;
      r = searchTree(tree->default_case, param);
      return r;
    }

    std::optional<Result> visitCaseTree(const PTR<CaseTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<ExpressionTree>(tree->conditions, param);
      if (r.has_value()) return r;
      r = searchTree(tree->body, param);
      return r;
    }

    std::optional<Result> visitForLoopTree(const PTR<ForLoopTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->target, param);
      if (r.has_value()) return r;
      r = searchTree(tree->start, param);
      if (r.has_value()) return r;
      r = searchTree(tree->end, param);
      if (r.has_value()) return r;
      r = searchTree(tree->step, param);
      if (r.has_value()) return r;
      r = searchTree(tree->body, param);
      return r;
    }

    std::optional<Result> visitForEachTree(const PTR<ForEachTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->source, param);
      if (r.has_value()) return r;
      r = searchTree(tree->var1, param);
      if (r.has_value()) return r;
      r = searchTree(tree->var2, param);
      if (r.has_value()) return r;
      r = searchTree(tree->body, param);
      return r;
    }

    std::optional<Result> visitWhileLoopTree(const PTR<WhileLoopTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->condition, param);
      if (r.has_value()) return r;
      r = searchTree(tree->body, param);
      return r;
    }

    std::optional<Result> visitLiteralTree(const PTR<LiteralTree>& tree, Param param) override {
      return std::nullopt;
    }

    std::optional<Result> visitArrayTypeTree(const PTR<ArrayTypeTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->expr, param);
      if (r.has_value()) return r;
      r = searchTree(tree->size, param);
      return r;
    }

    std::optional<Result> visitTemplateTypeTree(const PTR<TemplateTypeTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->class_name, param);
      if (r.has_value()) return r;
      r = searchTrees<ExpressionTree>(tree->args, param);
      return r;
    }

    std::optional<Result> visitIdentifierTree(const PTR<IdentifierTree>& tree, Param param) override {
      return std::nullopt;
    }

    std::optional<Result> visitFieldAccessTree(const PTR<FieldAccessTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->prefix, param);
      return r;
    }

    std::optional<Result> visitMethodInvokeTree(const PTR<MethodInvokeTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->method, param);
      if (r.has_value()) return r;
      r = searchTrees<ExpressionTree>(tree->args, param);
      return r;
    }

    std::optional<Result> visitNewClassTree(const PTR<NewClassTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->class_name, param);
      if (r.has_value()) return r;
      r = searchTrees<ExpressionTree>(tree->args, param);
      return r;
    }

    std::optional<Result> visitNewArrayTree(const PTR<NewArrayTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<ExpressionTree>(tree->elements, param);
      return r;
    }

    std::optional<Result> visitArrayAccessTree(const PTR<ArrayAccessTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->selected, param);
      if (r.has_value()) return r;
      r = searchTree(tree->index, param);
      return r;
    }

    std::optional<Result> visitBreakTree(const PTR<BreakTree>& tree, Param param) override {
      return std::nullopt;
    }

    std::optional<Result> visitContinueTree(const PTR<ContinueTree>& tree, Param param) override {
      return std::nullopt;
    }

    std::optional<Result> visitReturnTree(const PTR<ReturnTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->value, param);
      return r;
    }

    std::optional<Result> visitSubscribeEventTree(const PTR<SubscribeEventTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<IdentifierTree>(tree->vars, param);
      return r;
    }

    std::optional<Result> visitTypeCastTree(const PTR<TypeCastTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->expr, param);
      if (r.has_value()) return r;
      r = searchTree(tree->type_expr, param);
      return r;
    }

    std::optional<Result> visitAssignTree(const PTR<AssignTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->left, param);
      if (r.has_value()) return r;
      r = searchTree(tree->right, param);
      return r;
    }

    std::optional<Result> visitExchangeTree(const PTR<ExchangeTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->left, param);
      if (r.has_value()) return r;
      r = searchTree(tree->right, param);
      return r;
    }

    std::optional<Result> visitInstanceOfTree(const PTR<InstanceOfTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->expr, param);
      if (r.has_value()) return r;
      r = searchTree(tree->type_expr, param);
      return r;
    }

    std::optional<Result> visitMultiExprTree(const PTR<MultiExprTree>& tree, Param param) override {
      std::optional<Result> r = searchTrees<ExpressionTree>(tree->expressions, param);
      return r;
    }

    std::optional<Result> visitAwaitTree(const PTR<AwaitTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->expr, param);
      return r;
    }

    std::optional<Result> visitParensTree(const PTR<ParensTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->expr, param);
      return r;
    }

    std::optional<Result> visitUnaryTree(const PTR<UnaryTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->expr, param);
      return r;
    }

    std::optional<Result> visitBinaryTree(const PTR<BinaryTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->left, param);
      if (r.has_value()) return r;
      r = searchTree(tree->right, param);
      return r;
    }

    std::optional<Result> visitCodeTextTree(const PTR<CodeTextTree>& tree, Param param) override {
      return std::nullopt;
    }

    std::optional<Result> visitCodeRefTree(const PTR<CodeRefTree>& tree, Param param) override {
      std::optional<Result> r = searchTree(tree->expr, param);
      return r;
    }

    std::optional<Result> visitErroneousTree(const PTR<ErroneousTree>& tree, Param param) override {
      return std::nullopt;
    }
  };

  /// 语法树路径描述，支持从当前节点向上访问
  class TreePath : public std::enable_shared_from_this<TreePath> {
  public:
    explicit TreePath(const PTR<TreePath>& parent, const PTR<Tree>& leaf);
    PTR<Tree> getLeaf() const;
    PTR<CompilationUnitTree> getRootLeaf() const;
    PTR<TreePath> getParent() const;
    PTR<TreePath> rootPath() const;
    /// 遍历当前TreePath所处的作用域，会向上进行遍历
    /// @param consumer 遍历函数
    /// @return 返回当前TreePath所处的作用域，不会包含父级作用域(由于使用weak_ptr，父级作用域此时已经销毁)
    PTR<Scope> visitScope(std::function<void(const PTR<Scope>&)>&& consumer);

    static PTR<TreePath> getPath(const PTR<CompilationUnitTree>& root, const PTR<Tree>& target);
    static PTR<TreePath> getPath(const PTR<TreePath>& path, const PTR<Tree>& target);

  private:
    PTR<TreePath> parent_;
    PTR<Tree> leaf_;
    PTR<CompilationUnitTree> root_leaf_;
  };

  /// 按照语法树路径进行扫描(短路遍历)
  template <typename Result, typename Param>
  class TreePathSearcher : public TreeSearcher<Result, Param> {
  public:
    explicit TreePathSearcher(const PTR<TreePath>& path = nullptr) : path_(path) {
    }

    PTR<TreePath> getCurrentPath() const {
      return path_;
    }

    std::optional<Result> searchPath(const PTR<TreePath>& tree_path, Param param) {
      path_ = tree_path;
      return this->visitTree(path_->getLeaf(), param);
    }

    std::optional<Result> visitTree(const PTR<Tree>& tree, Param param) override {
      if (tree == nullptr) {
        return std::nullopt;
      }
      path_ = MAKE_PTR<TreePath>(path_, tree);
      return TreeSearcher<Result, Param>::visitTree(tree, param);
    }
  private:
    PTR<TreePath> path_{nullptr};
  };

  /// 按照语法树路径进行扫描(完整遍历)
  template <typename Result, typename Param>
  class TreePathScanner : public TreeScanner<Result, Param> {
  public:
    explicit TreePathScanner(const PTR<TreePath>& path = nullptr) : path_(path) {
    }

    PTR<TreePath> getCurrentPath() const {
      return path_;
    }

    Result scanPath(const PTR<TreePath>& tree_path, Param param) {
      path_ = tree_path;
      return scan(path_->getLeaf(), param);
    }

    Result scan(const PTR<Tree>& tree, Param param) override {
      if (tree == nullptr) {
        if constexpr (std::is_void_v<Result>) {
          return;
        } else {
          return {};
        }
      }
      path_ = MAKE_PTR<TreePath>(path_, tree);
      return TreeScanner<Result, Param>::scan(tree, param);
    }

  private:
    PTR<TreePath> path_{nullptr};
  };

  class TreePathFinder : public TreePathSearcher<PTR<TreePath>, const PTR<Tree>&> {
  public:
    std::optional<PTR<TreePath>> visitTree(const PTR<Tree>& tree, const PTR<Tree>& param) override {
      if (tree == param) {
        // 搜寻到匹配目标Tree的节点，不再继续扫描，直接返回TreePath
        return MAKE_PTR<TreePath>(getCurrentPath(), param);
      } else {
        return TreePathSearcher<PTR<TreePath>, const PTR<Tree>&>::visitTree(tree, param);
      }
    }
  };

  struct ScopeNode {
    const PTR<TreePath> target;
    PTR<MutableScope> scope;
    PTR<ClassTree> klass;
    PTR<MethodTree> method;

    ScopeNode(const PTR<TreePath>& target, const PTR<MutableScope>& scope,
      const PTR<ClassTree>& klass, const PTR<MethodTree>& method): target(target), scope(scope), klass(klass), method(method) {}

    static ScopeNode rootNode(const PTR<TreePath>& path) {
      return ScopeNode(path, nullptr, nullptr, nullptr);
    }

    ScopeNode dupForClass(const PTR<ClassTree>& class_tree) const {
      return ScopeNode(target, scope, class_tree, nullptr);
    }

    ScopeNode dupForMethod(const PTR<MethodTree>& method_tree) const {
      PTR<MutableScope> new_scope = MAKE_PTR<MutableScope>(scope, klass->symbol, method_tree->symbol);
      return ScopeNode(target, new_scope, klass, method_tree);
    }

    ScopeNode dupForBlock() const {
      PTR<MutableScope> new_scope = MAKE_PTR<MutableScope>(scope, klass->symbol, method->symbol);
      return ScopeNode(target, new_scope, klass, method);
    }
  };
  /// 基于TreePath扫描语法树所处的作用域
  class ScopeVisitor : public TreePathSearcher<PTR<MutableScope>, ScopeNode&> {
  public:
    template<typename Func, typename = std::enable_if_t<is_lambda_or_func<Func, void, const PTR<Scope>&>>>
    explicit ScopeVisitor(Func&& visitor): func_(std::forward<Func>(visitor)) {
    }

    PTR<Scope> visitScope(const PTR<TreePath>& target) {
      ScopeNode root_node = ScopeNode::rootNode(target);
      std::optional<PTR<MutableScope>> r = searchPath(target->rootPath(), root_node);
      if (r.has_value()) {
        return r.value();
      } else {
        return nullptr;
      }
    }

    std::optional<PTR<MutableScope>> visitTree(const PTR<Tree>& tree, ScopeNode& node) override {
      if (tree == node.target->getLeaf()) {
        PTR<Scope> scope = node.scope;
        while (scope != nullptr) {
          func_(scope);
          scope = scope->getEnclosingScope();
        }
        return node.scope;
      }
      return TreePathSearcher<PTR<MutableScope>, ScopeNode&>::visitTree(tree, node);
    }

    std::optional<PTR<MutableScope>> visitClassTree(const PTR<ClassTree>& tree, ScopeNode& node) override {
      ScopeNode class_node = node.dupForClass(tree);
      return TreePathSearcher<PTR<MutableScope>, ScopeNode&>::visitClassTree(tree, class_node);
    }

    std::optional<PTR<MutableScope>> visitMethodTree(const PTR<MethodTree>& tree, ScopeNode& node) override {
      ScopeNode method_node = node.dupForMethod(tree);
      return TreePathSearcher<PTR<MutableScope>, ScopeNode&>::visitMethodTree(tree, method_node);
    }

    std::optional<PTR<MutableScope>> visitEventTree(const PTR<EventTree>& tree, ScopeNode& node) override {
      ScopeNode event_node = node.dupForMethod(tree);
      return TreePathSearcher<PTR<MutableScope>, ScopeNode&>::visitEventTree(tree, event_node);
    }

    std::optional<PTR<MutableScope>> visitBlockTree(const PTR<BlockTree>& tree, ScopeNode& node) override {
      ScopeNode next_node = node.dupForBlock();
      return TreeSearcher<PTR<MutableScope>, ScopeNode&>::visitBlockTree(tree, next_node);
    }

    std::optional<PTR<MutableScope>> visitVariableTree(const PTR<VariableTree>& tree, ScopeNode& node) override {
      if (node.scope != nullptr) {
        node.scope->addSymbol(tree->symbol);
      }
      return std::nullopt;
    }
  private:
    std::function<void(const PTR<Scope>&)> func_;
  };
}

#endif